{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sigmoid(x): \n",
    "    return 1. / (1 + np.exp(-x))\n",
    "\n",
    "# createst uniform random array w/ values in [a,b) and shape args\n",
    "def rand_arr(a, b, *args): \n",
    "    np.random.seed(0)\n",
    "    return np.random.rand(*args) * (b - a) + a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "#所有的w，b，w_diff，b_diff\n",
    "class GruParam:\n",
    "    def __init__(self, x_dim, mem_cell_dim, inputs_dim):\n",
    "        self.x_dim=x_dim\n",
    "        self.mem_cell_dim = mem_cell_dim\n",
    "        concat_dim = x_dim + mem_cell_dim\n",
    "        self.w_r = rand_arr(-0.1,0.1,mem_cell_dim,concat_dim)\n",
    "        self.w_z = rand_arr(-0.1,0.1,mem_cell_dim,concat_dim)\n",
    "        self.w_h = rand_arr(-0.1,0.1,mem_cell_dim,concat_dim)\n",
    "        \n",
    "        self.u_r = rand_arr(-0.1,0.1,mem_cell_dim,concat_dim)\n",
    "        self.u_z = rand_arr(-0.1,0.1,mem_cell_dim,concat_dim)\n",
    "        self.u_h = rand_arr(-0.1,0.1,mem_cell_dim,concat_dim)\n",
    "        self.u_o = rand_arr(-0.1,0.1,mem_cell_dim,concat_dim)\n",
    "        \n",
    "        self.br = rand_arr(-0.1,0.1,mem_cell_dim)\n",
    "        self.bz = rand_arr(-0.1,0.1,mem_cell_dim)\n",
    "        self.bh = rand_arr(-0.1,0.1,mem_cell_dim)\n",
    "        self.wo = rand_arr(-0.1,0.1,mem_cell_dim)\n",
    "        \n",
    "        self.wr_diff = np.zeros((mem_cell_dim,concat_dim))\n",
    "        self.wz_diff = np.zeros((mem_cell_dim,concat_dim))\n",
    "        self.wh_diff = np.zeros((mem_cell_dim,concat_dim))\n",
    "        self.wo_diff = np.zeros((mem_cell_dim,concat_dim))\n",
    "        \n",
    "        self.br_diff = np.zeros(mem_cell_dim)\n",
    "        self.bz_diff = np.zeros(mem_cell_dim)\n",
    "        self.bh_diff = np.zeros(mem_cell_dim)\n",
    "        self.bo_diff = np.zeros(mem_cell_dim)\n",
    "        \n",
    "    def apply_diff(self,lr =1):\n",
    "        self.wr -= lr  *self.wr_diff\n",
    "        self.wz -= lr * self.wz_diff\n",
    "        self.wh -= lr * self.wh_diff\n",
    "        self.wo -= lr *   self.wo_diff\n",
    "        \n",
    "        self.wr_diff = np.zeros_like(self.wr)\n",
    "        self.wz_diff = np.zeros_like(self.wz)\n",
    "        self.wh_diff = np.zeros_like(self.wh)\n",
    "        self.wo_diff = np.zeros_like(self.wo)\n",
    "        self.br_diff = np.zeros_like(self.br)\n",
    "        self.bz_diff = np.zeros_like(self.bz)\n",
    "        self.bh_diff = np.zeros_like(self.bh)\n",
    "        self.bo_diff = np.zeros_like(self.bo)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 所有的中间值\n",
    "class GruState:\n",
    "    def __init__(self,x_dm, mem_cell_dim):\n",
    "        self.r = np.zeros(mem_cell_dim)\n",
    "        self.z = np.zeros(mem_cell_dim)\n",
    "        self.h = np.zeros(mem_cell_dim)\n",
    "        self.h_out = np.zeros(mem_cell_dim)\n",
    "        self.y_out = np.zeros(mem_cell_dim)\n",
    "        \n",
    "        self.diff_h = np.zeros_like(self.h)\n",
    "        self.diff_x = np.zeros_like(x_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GruNode:\n",
    "    def __init__(self,gru_param, gru_state):\n",
    "        self.state = gru_state\n",
    "        self.param = gru_param\n",
    "        self.x = None\n",
    "      \n",
    "    def forward(self, x, h_prev=None, flag=0):\n",
    "        if flag==0:\n",
    "            h_prev = np.zeros_like(self.state.h)\n",
    "        self.h_prev = h_prev\n",
    "        \n",
    "        self.state.r = np.dot(self.param.w_r,x) + np.dot(self.param.u_r, h_prev) + self.param.br\n",
    "        self.state.r = sigmoid(self.state.r)\n",
    "        self.state.z = np.dot(self.param.w_z,x) + np.dot(self.param.u_z, h_prev) + self.param.bz\n",
    "        self.state.z = sigmoid(self.state.z)\n",
    "        self.state.h = np.dot(self.param.w_h,x) + np.dot(self.param.u_h, self.state.r * h_prev) + self.param.bh\n",
    "        self.state.h = np.tanh(self.state.h)\n",
    "        self.state.h_next = h_prev * (1 - self.state.z) + self.state.z * self.state.h\n",
    "        self.y_pre = sigmoid(self.param.u_o @ self.state.h_next + self.state.bo)\n",
    "        \n",
    "        self.x = x\n",
    "     \n",
    "    def backward(self): \n",
    "        h_diff = h_next_diff * self.state.z * (1 - np.power(self.state.h,2))\n",
    "        r_diff = h_diff * self.param.u_h * h_prev * self.state.r * (1 - self.state.r)\n",
    "        z_diff = (h_next_diff * (self.state.h - h_prev)) * self.state.z * (1 - self.state.z)\n",
    "        h_prev_diff = self.param.u_r @ r_diff + self.param.u_z @ z_diff \n",
    "                        + self.param.u_h @ (h_diff * self.state.r) + h_next_diff * (1- self.state.z)\n",
    "        \n",
    "        wr_diff = x @ r_diff\n",
    "        ur_diff = h_prev @ r_diff\n",
    "        br_diff = r_diff\n",
    "        wz_diff = x * z_diff\n",
    "        uz_diff = h_prev * z_diff\n",
    "        bz_diff = z_diff\n",
    "        wh_diff = x @ h_diff\n",
    "        uh_diff = (self.state.r * h_prev) @ h_diff\n",
    "        bh_diff = h_diff\n",
    "        \n",
    "        \n",
    "    \n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "            \n",
    "        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
